function estimate( whiteningMatrix1, whiteningMatrix2, dewhiteningMatrix1, dewhiteningMatrix2, fname, p )
% estimate - algorithms for topographic, subspace, and standard ICA
%
% SYNTAX:
% estimate( whiteningMatrix, dewhiteningMatrix, fname, dims );
%
% NOTE: X passed as global!
%
% X                  preprocessed (whitened) sample vectors in columns
% whiteningMatrix    transformation from observation space to X-space
% dewhiteningMatrix  inverse transformation
% fname              name of file to write
%
% PARAMETERS COMMON TO ALL ALGORITHMS
%
% p.seed             random number generator seed
% p.write            iteration interval for writing results to disk
%
% ICA (with FastICA algorithm, tanh nonlinearity)
%
% p.model            'ica'
% p.algorithm        'fixed-point'
% p.components       number of ICA components to estimate
% 
% ISA (gradient descent with adaptive stepsize)
%
% p.model            'isa'
% p.algorithm        'gradient'
% p.groupsize        dimensionality of subspaces
% p.groups           number of independent subspaces to estimate
% p.stepsize         starting stepsize
% p.epsi             small positive constant
%
% TOPOGRAPHIC ICA (gradient descent with adaptive stepsize)
%
% p.model            'tica'
% p.algorithm        'gradient'
% p.xdim             columns in map
% p.ydim             rows in map
% p.maptype          'standard' or 'torus'
% p.neighborhood     'ones3by3' (only one implemented so far)
% p.stepsize         starting stepsize
% p.epsi             small positive constant
%

%-------------------------------------------------------------------
% PRELIMINARIES
%-------------------------------------------------------------------

% Print options
fprintf('You have selected the following options...\n');
p

% Neighborhood matrix: NB(i,j) = strength of unit j in neighb. of 
% unit i. In addition, we will create a matrix NBNZ, which gives 
% the positions of the non-zero entries in NB, to lower the 
% computational expenses.
fprintf('Generating neighborhood matrix...\n');
[NBNZ,NB] = GenerateNB( p );
  
% Take the data from the global variable
global X1 X2;
N1 = size(X1,2);
N2 = size(X2,2);

%-------------------------------------------------------------------
% SETTING UP THE STARTING POINT...
%-------------------------------------------------------------------

% Initialize the random number generator.
rand('seed',p.seed);

% Take random initial vectors...
B1 = randn(size(X1,1),p.xdim*p.ydim);  
B2 = randn(size(X2,1),p.xdim*p.ydim);  
  
% ...and decorrelate (=orthogonalize in whitened space)

if strcmp(p.mode,'cpu')
    B1 = B1*real((B1'*B1)^(-0.5));
    B2 = B2*real((B2'*B2)^(-0.5));
elseif strcmp(p.mode,'gpu') | strcmp(p.mode,'GPU')
    B1 = double(matCudaXTICA('decorrelate', single(B1), int32(p.maxDecoIter), single(p.decoEpsi)));
    B2 = double(matCudaXTICA('decorrelate', single(B2), int32(p.maxDecoIter), single(p.decoEpsi)));
end

n1 = size(B1,2);
n2 = size(B2,2);

%-------------------------------------------------------------
% Normalize 
%-------------------------------------------------------------
  
if strcmp(p.mode, 'cpu')
    Bsz = size(B1, 2);
    for i = 1:Bsz
        B1(:,i) = B1(:,i)/norm(B1(:,i),2);
        B2(:,i) = B2(:,i)/norm(B2(:,i),2);
    end
elseif strcmp(p.mode, 'gpu') | strcmp(p.mode, 'GPU')
    B1 = double(matCudaXTICA('normalize', single(B1)));
    B2 = double(matCudaXTICA('normalize', single(B2)));
end
  
%-------------------------------------------------------------------
% START THE ITERATION...
%-------------------------------------------------------------------

% Print the time when started (and save along with parameters).
c=clock;
if c(5)<10, timestarted = ['Started at: ' int2str(c(4)) ':0' int2str(c(5))];
else timestarted = ['Started at: ' int2str(c(4)) ':' int2str(c(5))];
end
fprintf([timestarted '\n']);
p.timestarted = timestarted;

% Initialize iteration counter
iter=0;

% Use starting stepsize for gradient methods
stepsize1 = p.stepsize; 
stepsize2 = p.stepsize; 
obj1 = [];
obj2 = [];
objiter1 = [];
objiter2 = [];

% Loop forever, writing result periodically
while iter < p.maxIter  
    % Increment iteration counter
    iter = iter+1;  
    fprintf('(%d)',iter);

    if strcmp(p.mode,'GPU')
        [B1, B2, stepsize1, stepsize2, objective1, objective2] = matCudaXTICA('iteraction', single(B1), single(B2), single(X1), single(X2), single(NB), int32(p.gpuIter), int32(p.update), single(stepsize1), single(stepsize2), p.orthoMode, single(p.localOrthoGamma), single(p.globalOrthoGamma), single(p.orthoEpsi), single(p.epsi), int32(p.maxDecoIter), single(p.decoEpsi));      
        
        c=clock;
        if c(5)<10, itertime = ['One iteration finished at: ' int2str(c(4)) ':0' int2str(c(5))];
        else itertime = ['One iteration finished at: ' int2str(c(4)) ':' int2str(c(5))];
        end
        
        fprintf('\n%s\n',itertime);
        
        B1 = double(B1);
        B2 = double(B2);
        stepsize1 = double(stepsize1);
        stepsize2 = double(stepsize2);
        objective1 = double(objective1);
        objective2 = double(objective2);

        fprintf('Stepsize 1 now: %.4f\n',stepsize1);
        fprintf('Stepsize 2 now: %.4f\n',stepsize2);
        fprintf('Objective 1 value: %.6f\n',objective1);      
        fprintf('Objective 2 value: %.6f\n',objective2);   
               
        objiter1 = [objiter1 p.gpuIter*iter];
        obj1 = [obj1 objective1];
        objiter2 = [objiter2 p.gpuIter*iter];
        obj2 = [obj2 objective2];

        if rem(iter,p.write)==0 | iter==1

            A1 = dewhiteningMatrix1 * B1;
            A2 = dewhiteningMatrix2 * B2;
            W1 = B1' * whiteningMatrix1;
            W2 = B2' * whiteningMatrix2;

            fprintf(['Writing file: ' fname '...']);
            
            eval(['save ' fname ' W1 W2 A1 A2 p iter obj1 obj2 objiter1 objiter2 ']);
            
            fprintf(' Done!\n');      
        end
      
      
    else  

        %-------------------------------------------------------------
        % TICA gradient
        %-------------------------------------------------------------  

        if strcmp(p.mode,'cpu')
            % Calculate linear filter responses and their squares
            U11 = B1'*X1; Usq11 = U11.^2;
            U12 = B1'*X2; Usq12 = U12.^2;
            U21 = B2'*X1; Usq21 = U21.^2;
            U22 = B2'*X2; Usq22 = U22.^2;

            % Calculate local energies
            for ind=1:n
                E11(ind,:) = NB(ind,NBNZ{ind}) * Usq11(NBNZ{ind},:);
                E12(ind,:) = NB(ind,NBNZ{ind}) * Usq12(NBNZ{ind},:);
                E21(ind,:) = NB(ind,NBNZ{ind}) * Usq21(NBNZ{ind},:);
                E22(ind,:) = NB(ind,NBNZ{ind}) * Usq22(NBNZ{ind},:);
            end

            % Take nonlinearity
            g11 = -((p.epsi + E11).^(-0.5));
            g12 = -((p.epsi + E12).^(-0.5));
            g21 = -((p.epsi + E21).^(-0.5));
            g22 = -((p.epsi + E22).^(-0.5));

            % Calculate convolution with neighborhood
            for ind=1:n
                F11(ind,:) = NB(ind,NBNZ{ind}) * g11(NBNZ{ind},:);
                F12(ind,:) = NB(ind,NBNZ{ind}) * g12(NBNZ{ind},:);
                F21(ind,:) = NB(ind,NBNZ{ind}) * g21(NBNZ{ind},:);
                F22(ind,:) = NB(ind,NBNZ{ind}) * g22(NBNZ{ind},:);
            end

            % This is the total gradient
            dB11 = X1*(U11.*F11)'/N1;
            dB12 = X2*(U12.*F12)'/N2;
            dB21 = X1*(U21.*F21)'/N1;
            dB22 = X2*(U22.*F22)'/N2;

            dB1 = dB11-dB12;
            dB2 = dB22-dB21;

            if strcmp(p.orthoMode,'orthogonalize')
                BtB11 = B1'*B1;
                BtB12 = B1'*B2;
                BtB21 = B2'*B1;
                BtB22 = B2'*B2;
                Ou11 = -2*BtB11;
                Ou12 = -2*BtB12;
                Ou21 = -2*BtB21;
                Ou22 = -2*BtB22;
                Ol11 = 1-BtB11.*BtB11;
                Ol12 = 1-BtB12.*BtB12;
                Ol21 = 1-BtB21.*BtB21;
                Ol22 = 1-BtB22.*BtB22;
                BtBsz11 = size(BtB11, 2);
                BtBsz12 = size(BtB12, 2);
                BtBsz21 = size(BtB21, 2);
                BtBsz22 = size(BtB22, 2);

                for i = 1:BtBsz11
                    Ou11(i,i) = 0;
                    Ol11(i,i) = 1;
                end

                for i = 1:BtBsz22
                    Ou22(i,i) = 0;
                    Ol22(i,i) = 1;
                end

                OO11 = Ou11./Ol11;
                OO12 = Ou12./Ol12;
                OO21 = Ou21./Ol21;
                OO22 = Ou22./Ol22;

                for i = 1:BtBsz11
                    for j = 1:BtBsz11
                        dB1(:,i) = dB1(:, i) + OO11(i,j)*B1(:,j);
                    end
                end
                for i = 1:BtBsz12
                    for j = 1:BtBsz12
                        dB1(:,i) = dB1(:, i) + OO12(i,j)*B2(:,j);
                    end
                end
                for i = 1:BtBsz21
                    for j = 1:BtBsz21
                        dB2(:,i) = dB2(:, i) + OO21(i,j)*B1(:,j);
                    end
                end
                for i = 1:BtBsz22
                    for j = 1:BtBsz22
                        dB2(:,i) = dB2(:, i) + OO22(i,j)*B2(:,j);
                    end
                end

            end

        elseif strcmp(p.mode,'gpu')
            dB1 = double(matCudaXTICA('gradient', single(B1), single(X1), single(X2), single(NB), single(p.epsi)));       
            dB2 = double(matCudaXTICA('gradient', single(B2), single(X2), single(X1), single(NB), single(p.epsi)));       

            if strcmp(p.orthoMode,'orthogonalize')
                dB1 = double(matCudaXTICA('orthogonalize1', single(dB1), single(B1), single(p.localOrthoGamma), single(p.orthoEpsi)));       
                dB1 = double(matCudaXTICA('orthogonalize2', single(dB1), single(B1), single(B2), single(p.globalOrthoGamma), single(p.orthoEpsi)));       

                dB2 = double(matCudaXTICA('orthogonalize1', single(dB2), single(B2), single(p.localOrthoGamma), single(p.orthoEpsi)));       
                dB2 = double(matCudaXTICA('orthogonalize2', single(dB2), single(B2), single(B1), single(p.globalOrthoGamma), single(p.orthoEpsi)));
            end
        end    

        %-------------------------------------------------------------
        % ADAPT STEPSIZE FOR GRADIENT ALGORITHMS
        %-------------------------------------------------------------  

        if strcmp(p.mode,'cpu') 
            % Perform this adaptation only every 5 steps
            if rem(iter,p.update)==0 | iter==1

                % Take different length steps
                Bc{1} = B1 + 0.5*stepsize1*dB1;
                Bc{2} = B1 + 1.0*stepsize1*dB1;
                Bc{3} = B1 + 2.0*stepsize1*dB1;

                % Orthogonalize each one
                for i=1:3, Bc{i} = Bc{i}*real((Bc{i}'*Bc{i})^(-0.5)); end

                % Calculate objective values in each case
                for i=1:3
                    Usq = (Bc{i}'*X).^2;
                    for ind=1:n, E(ind,:)= NB(ind,NBNZ{ind}) * Usq(NBNZ{ind},:); end
                    objective(i) = mean(mean(sqrt(p.epsi+E)));
                end

                % Compare objective values, pick step giving minimum
                if objective(1)<objective(2) & objective(1)<objective(3)
                    % Take shorter step
                    stepsize1 = stepsize1/2;
                    fprintf('Stepsize now: %.4f\n',stepsize);
                    obj1 = [obj1 objective(1)];
                elseif objective(1)>objective(3) & objective(2)>objective(3)
                    % Take longer step
                    stepsize1 = stepsize1*2;
                    fprintf('Stepsize now: %.4f\n',stepsize);
                    obj1 = [obj1 objective(3)];
                else
                    % Don't change step
                    obj1 = [obj1 objective(2)];
                end

                objiter = [objiter obj(end)];
                
                fprintf('\nStepsize 1 now: %.4f\n',stepsize1);
                fprintf('Stepsize 2 now: %.4f\n',stepsize2);
                fprintf('Objective 1 value: %.6f\n',objective1);            
                fprintf('Objective 2 value: %.6f\n',objective2); 
            end

            B1 = B1 + stepsize*dB1;
            
        elseif strcmp(p.mode,'gpu')  
            % Perform this adaptation only every 5 steps
            if rem(iter,p.update)==0 || iter==1
                [stepsize1, objective1] = matCudaXTICA('adaption', single(B1), single(dB1), single(X1), single(NB), single(stepsize1), int32(100), single(1e-3));
                [stepsize2, objective2] = matCudaXTICA('adaption', single(B2), single(dB2), single(X2), single(NB), single(stepsize2), int32(100), single(1e-3));
                objiter1 = [objiter1 iter];
                obj1 = [obj1 objective1];
                objiter2 = [objiter2 iter];
                obj2 = [obj2 objective2];

                stepsize1 = double(stepsize1);
                stepsize2 = double(stepsize2);
                obj1 = double(obj1);       
                objiter2 = double(objiter2);       
                fprintf('\nStepsize 1 now: %.4f\n',stepsize1);
                fprintf('Stepsize 2 now: %.4f\n',stepsize2);
                fprintf('Objective 1 value: %.6f\n',objective1);            
                fprintf('Objective 2 value: %.6f\n',objective2);                       
            end
            B1 = double(matCudaXTICA('update', single(B1), single(dB1), single(stepsize1)));            
        end

        %-------------------------------------------------------------
        % Ortogonalize (equal to decorrelation since we are 
        % in whitened space)
        %-------------------------------------------------------------

        if strcmp(p.orthoMode, 'decorrelate')
            if strcmp(p.mode, 'cpu')
                B = B*real((B'*B)^(-0.5));
            elseif strcmp(p.mode, 'gpu')
                B1 = double(matCudaXTICA('decorrelate', single(B1), int32(p.maxDecoIter), single(p.decoEpsi)));
                B2 = double(matCudaXTICA('decorrelate', single(B2), int32(p.maxDecoIter), single(p.decoEpsi)));
            end
        end

        %-------------------------------------------------------------
        % Normalize 
        %-------------------------------------------------------------

        if strcmp(p.mode, 'cpu')
            Bsz = size(B, 2);
            for i = 1:Bsz
                B(:,i) = B(:,i)/norm(B(:,i),2);
            end
        elseif strcmp(p.mode, 'gpu')
            B1 = double(matCudaXTICA('normalize', single(B1)));
            B2 = double(matCudaXTICA('normalize', single(B2)));
        end

        %-------------------------------------------------------------
        % Write the results to disk
        %-------------------------------------------------------------

        if rem(iter,p.write)==0 | iter==1

            A1 = dewhiteningMatrix1 * B1;
            A2 = dewhiteningMatrix2 * B2;
            W1 = B1' * whiteningMatrix1;
            W2 = B2' * whiteningMatrix2;

            fprintf(['Writing file: ' fname '...']);
            
            eval(['save ' fname ' W1 W2 A1 A2 B1 B2 p iter obj1 obj2 objiter1 objiter2']);
            
            fprintf(' Done!\n');      

        end
    end
end

% We never get here...
return;




%-----------------------------------------------------------------
% GenerateNB - generates the neighborhood matrix for TICA
%-----------------------------------------------------------------
function [NBNZ,NB] = GenerateNB( p )

% This will hold the neighborhood function entries
NB = zeros(p.xdim*p.ydim*[1 1]);

% This is currently the only implemented neighborhood
if strcmp(p.neighborhood,'ones3by3')==0
  error('No such neighborhood allowed!');
end

% Step through nodes one at a time to build the matrix
ind = 0;
for y=1:p.ydim
  for x=1:p.xdim
    
    ind = ind+1;
    
    if strcmp(p.maptype,'identical')
        NB( ind, ind )=1;
        
    else
        % Rectangular neighbors
        [xn,yn] = meshgrid( (x-1):(x+1), (y-1):(y+1) );
        xn = reshape(xn,[1 9]);
        yn = reshape(yn,[1 9]);

        if strcmp(p.maptype,'torus')

          % Cycle round
          i = find(yn<1); yn(i)=yn(i)+p.ydim;
          i = find(yn>p.ydim); yn(i)=yn(i)-p.ydim;
          i = find(xn<1); xn(i)=xn(i)+p.xdim;
          i = find(xn>p.xdim); xn(i)=xn(i)-p.xdim;

        elseif strcmp(p.maptype,'standard')

          % Take only valid nodes
          i = find(yn>=1 & yn<=p.ydim & xn>=1 & xn<=p.xdim);
          xn = xn(i);
          yn = yn(i);

        else
          error('No such map type!');
        end
        % Set neighborhood
        NB( ind, (yn-1)*p.xdim + xn )=1;
    end
    
  end
end

% For each unit, calculate the non-zero columns!
for i=1:p.xdim*p.ydim
  NBNZ{i} = find(NB(i,:));
end

return;

